So, the Finalize step only considers the segmenter values for state[0]? I'm not sure I understand why.

See [1], the Finalize is following a Merge call, by when the states of all threads have been merged into state[0].

But this raises an interesting question that I only realized now: the proposed thread-local segmenter_valuess are never merged, how dare I only append segmenter_values in state[0]? This seems wrong but turns out to be OK because of some implicit assumptions quite far away:

1. If the segmenter_values_ is not empty, then at least one segment key is specified, then we are executing single-threaded and only "state[0]" exists. This is detailed in my answer to the other comment.
2. Else, then there are no segment keys and we are appending nothing.

Even though the fix actually works, it "feels" more weird than before - a shared segmenter_values_ seems to be, at least conceptually, more reasonable (though it has race and isn't really correct in the context of multi-threading). This makes me hesitate about the current fix. I think it all boils down to the "abstraction leak" in the original design:

1. The segmenter abstraction seems to be designed to be independent of the assumption that "no multi-threading for segmented cases" (and honestly I am personally quite fond of this design).
2. The existence of segmenter_values_ strongly depends on the assumption of single-threaded.

Now I will try to fix the race in another way more independent of this abstraction leak, leaving the later to be addressed in the future.

[1]

Yes, that sounds reasonable. If we wanted to have multi-threaded segmented group-by, I suppose it would need a preparatory step to rechunk the input along segment boundaries?

It can be the case, but may not be that useful. Because IMO part of the point of segmented aggregation is to emit the result in a streaming fashion, that is, as soon as a segment is concluded to be "close", we are sure that the current aggregation result is a valid partial result thus can be output to downstream. Rechunking the input would require all input batches to be accumulated already.

To multi-threading the segmented aggregation, I would imagine the batches to be already partitioned by (and sorted by, of course, this is already implied by the current single-threaded impl) segment keys and distributed to specific threads (# partition == # thread). This can be achieved by a special source node or a "shuffle" node.

This way each thread of the aggregate can process all rows belonging to a specific segment. This would require some modification to the current aggregate node such as not merging other thread states, or just have a brand new "partitioned aggregate" node.

Anyway, it's not trivial and can be quite restrictive to use.

BTW, I've redone the fix in a way that is less weird.